CUDA support in blocks - #1599
Merged
Merged
Conversation
Everything downstream of an AVFrame's owner took `UniqueAVFrame&` or `const UniqueAVFrame&`. That's a constraint on the caller's storage rather than a statement about what the function does, and it has two costs. A non-const `UniqueAVFrame&` lets a callee take ownership of the caller's frame. Both CUDA interfaces did: BetaCudaDeviceInterface moved out of it, CudaDeviceInterface reassigned it. That is invisible under SingleStreamDecoder, whose frame is a loop local that dies right after conversion, but the building-block ops own their frame in a handle that outlives the call, so the frame gets freed twice. `const UniqueAVFrame&` doesn't allow the steal, but it still forces anyone holding a plain AVFrame* -- which is what the ops' tensor handle really is -- to manufacture a unique_ptr just to make the call, and manufacturing a second owner for an already-owned object is its own bug factory. So: functions that only look at a frame now take `const AVFrame&`, and the one that writes to it takes `AVFrame&`. Ownership stays with whoever actually owns the frame. Producers (receive_frame) keep `UniqueAVFrame&` because they really do hand back ownership. With that, the ops layer needs no ownership sleight-of-hand: wrap_pointer_to_tensor() gains a deleter parameter, so the one generic handle covers Demuxer/PacketDecoder/ColorConverter and the FFmpeg types too, and both bespoke wrap_*_pointer_to_tensor() functions go away. Demuxer::next_packet() returns UniqueAVPacket rather than a raw pointer plus a comment telling the caller to free it. The encoder is left alone: it owns and mutates its frames, and it uses a null frame as the flush signal, so a reference is the wrong shape there.
Two things only show up when building against the older FFmpeg headers. get_num_channels() was patching av_frame.channel_layout when FFmpeg 4 left it unset, so it was a mutator wearing an observer's name, and the layout fix-up was a side effect that swresample setup silently relied on. Pull the fix-up into get_channel_layout() and call that from the two places that actually need a layout; get_num_channels() just counts. swr_alloc_set_opts2() only became const-correct in FFmpeg 6 (libswresample 4.12). Cast for the older headers, which don't modify the layout either.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/meta-pytorch/torchcodec/1599
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This enables our NVDEC CUDA interface on the "Blocks" APIs. Lots of things are still not supported but the basic functionality works and is on par to CPU.
The original implementation was vibe-coded in #1592, and I re-implemented everything here from scratch because these aren't trivial changes and I really needed to understand them.
A few key points:
PacketDecoderin a newmake_frame_standalone()method before it is returned (and potentially later color-converted). That is needed to satisfy an assumption of the NVDEC decoding loop: we always unmap just before we map a new frame, and for that to be correct, the frame that we unmap must not be needed anymore. InSingleStreamDecodermode, we rely on the fact that this frame (that we unmap) is either discarded (we don't care about it) or color-converted (and thus copied), which makes it unmap-able. Here, we cannot assume that anymore, so we enforce a manual copy.convert_av_frame_to_frame_output()) was previously relying on the interface's state - a state that came from video decoding parts, which aren't present in its color-convert-only mode. This PR makesconvert_av_frame_to_frame_output()rely on info that doesn't come from that state. Particularly, part of that state is now attached to theAVFramevia theStandAloneFrameAttachedDatastruct.CudaContextGuardand associated comment.